(* Copyright Per Lindgren 2015, see the file "LICENSE" *)
(* for the full license governing this code.           *)

(* FUN-4SURE/Well *)


open Common 
open AST
open ST
open Env
open List
open Options

let tmatchl err t tl  =
  if List.mem t tl then t else 
    raise (RtfmError("Type error, got " ^ 
    st_of_ptype t ^ " expected " ^ myconcat " {" (List.map st_of_ptype tl) ^ "} in " ^ err )) 

let tmatch err e_t t  =
  if t == e_t then e_t else 
    raise (RtfmError("Type error, got " ^ 
    st_of_ptype t ^ " expected " ^ st_of_ptype e_t ^ " in " ^ err))
                                            
let rec unify env err e1 e2 =  
  let t1 = type_of_expr env (st_of_exp e1) e1 in
  let t2 = type_of_expr env (st_of_exp e2) e2 in
  tmatch err t1 t2
  
and type_of_op err e1 e2 env = function
  | OpAnd  -> tmatch  err (unify env err e1 e2) Bool
  | OpOr   -> tmatch  err (unify env err e1 e2) Bool
  | OpExor -> tmatch  err (unify env err e1 e2) Bool
  | OpMul  -> tmatch  err (unify env err e1 e2) Int
  | OpDiv  -> tmatch  err (unify env err e1 e2) Int
  | OpMod  -> tmatch  err (unify env err e1 e2) Int
  | OpAdd  -> tmatch  err (unify env err e1 e2) Int
  | OpSub  -> tmatch  err (unify env err e1 e2) Int 
  | OpEq   -> tmatchl err (unify env err e1 e2) [Int; Char; Bool] 
  | OpNeq  -> tmatchl err (unify env err e1 e2) [Int; Char; Bool]
  | OpGt   -> tmatchl err (unify env err e1 e2) [Int; Char]
  | OpGEq  -> tmatchl err (unify env err e1 e2) [Int; Char]
  | OpLt   -> tmatchl err (unify env err e1 e2) [Int; Char]
  | OpLEq  -> tmatchl err (unify env err e1 e2) [Int; Char]
  
and type_of_expr env err = function
  | OpExp   (op, e1, e2)  -> type_of_op err e1 e2 env op
  | IdExp   (idl)         -> t_of_id idl env 
  | IntExp  (i)           -> Int
  | CharExp (c)           -> Char
  | BoolExp (b)           -> Bool
  | UnOpExp (UNot, e)     -> tmatch err Bool (type_of_expr env err e)
  | UnOpExp (UMinus, e)   -> tmatch err Int (type_of_expr env err e)
    
let check_exp env e =
  let err = st_of_exp e in
  type_of_expr env err e  
  
let check_exp_to env t e =
  let err = st_of_exp e in
  let t_got = type_of_expr env err e in
  match (t, e) with
  | (Bool, IntExp 0) -> Bool
  | (Bool, IntExp 1) -> Bool
  | _                ->
  tmatch err t t_got                  
                 
let str_of_ids l =
  String.concat ", " (map st_of_id l)

let rec check_dup err = function
  | []     -> ()
  | e :: l -> 
    if mem e l then
     raise (RtfmError("Duplicate defintions of " ^ st_of_id e ^ ", " ^ err ))
    else check_dup err l
    
let well p = 
  
  let well_fb = function
    | FB (fb_id, fb_if, fb_v, fb_ecc, fb_alg) -> 
    (* error message *)
    let err = ", function block " ^ st_of_id fb_id ^ "." in
        
    (* event inputs*)  
    let eil = eil_of_if fb_if in 
    let ei_idl = map id_of_e eil in 
    
    (* event outputs *)
    let eol = eol_of_if fb_if in
    let oe_idl = map id_of_e eol in
    
    (* variable inputs *)
    let vil    = vil_of_if fb_if in 
    let vi_idl = map id_of_va vil in 

    (* variable inputs *)
    let vol    = vol_of_if fb_if in
    let vo_idl = map id_of_va vol in

    (* variable internals *)
    let vl     = vl_of_v fb_v in
    let v_idl  = map id_of_va vl in
    
    (* event input with constructs *)
    let eiwvl  = List.concat (filt_opt wio_of_e eil) in
    debug ("WITH input varaibles " ^ st_of_idl eiwvl); 
    iter (lookup vi_idl (", WITH input variable" ^ err)) eiwvl; 
    
    (* event output with constructs *)
    let eowvl  = List.concat (filt_opt wio_of_e eol) in
    debug ("WITH output variables " ^ st_of_idl eowvl); 
    iter (lookup vo_idl (", WITH output variable" ^ err)) eowvl; 
    
    (* algorithm identifiers *)
    let a_idl  = map id_of_alg fb_alg in
    
    verb ("Checking " ^ st_of_id fb_id );
    
    (* check all identifiers for dublicates *)
    let all_id = concat [ei_idl; oe_idl; vi_idl; vo_idl; v_idl; a_idl]; in
    check_dup err all_id;
    
    (* environment, id -> pType *)
    let env = map v_of_va (concat [vil; vol; vl]) in
        
    (* check ecc *)  
    match fb_ecc with
    | Some (FbEcc (s, _, tr)) -> ( 
      let err_ecc = "in ECC" ^ err in
      (* check duplicate states *)
      let ec_s = map sid_of_st s in
      check_dup err_ecc ec_s;
      
      (* check source and destination *)
      let tr_sd = map sd_of_tr tr in
      
      (* check source *)
      let tr_s = map fst tr_sd in
      iter (lookup ec_s (", source edge " ^ err_ecc)) tr_s;  
      
      (* check destination *)
      let tr_d = map snd tr_sd in
      iter (lookup ec_s (", destination edge " ^ err_ecc)) tr_d;  
      
      (* check state actions *)
      
      (* all actions *)
      let actl = concat (map act_of_st s) in
      
      (* all algorithm identifiers for all actions *) 
      let alg_idl = filt_opt aid_of_act actl in
      debug ("a_idl = " ^ String.concat ", " (map st_of_id a_idl));
      iter (lookup a_idl (", state action : algorithm identifier " ^ err_ecc)) alg_idl;
      
      (* all output events for all actions *)
      let ev_idl = concat (map eid_of_act actl) in
      iter (lookup oe_idl (", state action : output event identifier " ^ err_ecc)) ev_idl;
                      
      (* check all transition guard input events *)
      let ev_idl = filt_opt eid_of_tr tr in
      iter (lookup ei_idl (", state transition guard : input event identifier " ^ err_ecc)) ev_idl;
      
      (* check transition guard expressions *)
      let tr_g = map gu_of_tr tr in
      debug ("tr_g " ^ nl ^ String.concat nl (map st_of_gu tr_g));
      let tr_e = map exp_of_gu tr_g in
      debug ("tr_e " ^ nl ^ String.concat nl (map st_of_exp tr_e));
      
      verb ("Variable environment : {" ^  st_of_t_env env ^ "}"); 
      try 
        (* type check each expression, against Bool *)
        let _ = map (check_exp_to env Bool) tr_e in
 
        verb ("Typecheck algorithms");      
        (* type check each algorithm *)
        let check_alg env =
          let check_stmt env = function
            | StmtAssign (lid, exp) -> let _ = check_exp_to env (t_of_id lid env) exp in ()
            | StmtAssert (exp)      -> let _ = check_exp_to env Bool in ()
            
          in function
            | FbAlg (alg_id, _, vl, stmtl) ->
              (* extended environment with local variables *)
              let l_env = map v_of_va vl in
              let env = List.append l_env env in
              check_dup ("in " ^ st_of_id alg_id) (map fst env);
              iter (check_stmt env) stmtl 
             
        in
        iter (check_alg env) fb_alg ;  
         
        ()
      with
      | RtfmError msg -> raise (RtfmError (msg ^ err))
      | _             -> raise (RtfmError "Internal compiler error in Well formed!")       
    ) 
    | _ -> ()
    
    
  in
  verb ("Well formeness and typecking");
  map well_fb p 
  
 